import source
import ai.utils.utils_file as utils_file
import tqdm
from ai.thread.my_thread import MyThreadPool
from ai.utils import utils_data


def filter_dic(dic: dict, res: dict, min_sec=5, index: int = 0):
    """"""
    total_num = len(dic.items())
    if index == 0:
        tqdm_object = tqdm.tqdm(dic.items(), total=total_num)
    else:
        tqdm_object = dic.items()
    for id, dic_value in tqdm_object:
        wav_path = dic_value['wav']
        wav_form_num, sr = utils_data.get_sample_count(wav_path)
        if wav_form_num > min_sec * sr:
            res[id] = dic_value


def filter_file(json_file: str, min_sec=5, ):
    """"""
    res = {}
    dic = utils_file.load_json_file(json_file)
    total_keys = list(dic.keys())
    total_num = len(total_keys)
    my_thread_pool = MyThreadPool()
    thread_num = 50
    single_dic_lens = total_num // thread_num
    for i in range(0, total_num, single_dic_lens):
        dic_small = {}
        dic_small_keys = total_keys[i:i + single_dic_lens]
        for key in dic_small_keys:
            dic_small[key] = dic[key]
        my_thread_pool.add_thread(filter_dic, [dic_small, res, min_sec, i])
    my_thread_pool.start()
    output_path = json_file.strip().split('.json')[0] + '_filtered.json'
    utils_file.write_dict_to_json(res, output_path)


def merge_dic_for_json(json_file1: str, json_file2: str):
    res_dic = {}
    dic1 = utils_file.load_json_file(json_file1)
    dic2 = utils_file.load_json_file(json_file2)
    for key, value in dic1.items():
        res_dic[key] = value
    for key, value in dic2.items():
        if key not in res_dic:
            res_dic[key] = value
        else:
            continue
    utils_file.write_dict_to_json(res_dic, './output/accent_all.json')


def get_num_for_regions_for_json(source_json_file: str, output_json_file: str = './output/num_for_regions.json'):
    """
    首先得到每种口音的数据量是多少
    """
    count = {}
    dic = utils_file.load_json_file(source_json_file)
    for key, value in dic.items():
        if value['spk_id'] not in count:
            count[value['spk_id']] = 1
        else:
            count[value['spk_id']] += 1
    utils_file.write_dict_to_json(count, output_json_file)
    print(count)


def get_same_for_region_json(source_json_file: str, num_lens=10000,
                             output_json_file: str = './output/accent_same.json'):
    """
    对每种口音指定数量
    """
    dic = utils_file.load_json_file(source_json_file)
    # 定义一个空字典new_dic，用来存储新的字典
    new_dic = {}
    # 定义一个空列表kinds，用来存储所有的kind值
    kinds = []
    # 遍历dic中的每个键值对
    for key, value in dic.items():
        # 获取value中的kind值
        kind = value["spk_id"]
        # 如果kind不在kinds中，就将其添加到kinds中
        if kind not in kinds:
            kinds.append(kind)
    print(kinds)
    # 遍历kinds中的每个元素
    for kind in kinds:
        # 定义一个计数器count，用来记录每种kind的数量
        count = 0
        # 遍历dic中的每个键值对
        for key, value in dic.items():
            # 如果value中的kind值等于当前的kind值，并且count小于10，就将键值对添加到new_dic中，并将count加一
            if value["spk_id"] == kind and count < num_lens:
                new_dic[key] = value
                count += 1
            # 如果count等于10，就跳出循环
            if count == num_lens:
                break
    utils_file.write_dict_to_json(new_dic, output_json_file)


def get_same_for_region_json_for_new_dataset(source_json_file: str,
                                             output_json_file: str = './output/accent_filter.json'):
    """
    对每种口音指定数量
    """
    full_count = 30000
    dic = utils_file.load_json_file(source_json_file)
    # 定义一个空字典new_dic，用来存储新的字典
    new_dic = {}
    # 定义一个空列表kinds，用来存储所有的kind值
    kinds = []
    # 遍历dic中的每个键值对
    for key, value in dic.items():
        # 获取value中的kind值
        kind = value["spk_id"]
        # 如果kind不在kinds中，就将其添加到kinds中
        if kind not in kinds:
            kinds.append(kind)
    print(kinds)
    # 遍历kinds中的每个元素
    for kind in kinds:
        # 定义一个计数器count，用来记录每种kind的数量
        count = 0
        # 遍历dic中的每个键值对
        for key, value in dic.items():
            # 如果value中的kind值等于当前的kind值，并且count小于10，就将键值对添加到new_dic中，并将count加一
            if value["spk_id"] == kind and kind == 'Mandarin' and count < full_count:
                new_dic[key] = value
                count += 1
                # 如果count等于10，就跳出循环
                if count == full_count:
                    break
            if value["spk_id"] == kind:
                new_dic[key] = value
    utils_file.write_dict_to_json(new_dic, output_json_file)


if __name__ == '__main__':
    """"""
    # merge_dic_for_json('./output/accent_train.json', './output/accent_dev.json')
    # get_same_for_region_json('./output/accent_all.json.json')
    # get_num_for_regions_for_json('./output/accent_same.json', './output/num_for_regions4.json')
    get_same_for_region_json_for_new_dataset('./output/accent_train.json', './output/accent_filter.json')
